import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Subset
from sklearn.model_selection import train_test_split
import os
import logging
from prepare_data import DACDataset
from deep_fm import DeepFM
from train import train_model

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler()]
)


def main():
    # 配置参数
    batch_size = 100
    embedding_dim = 8
    hidden_dims = [512, 256, 128, 64]
    lr = 0.001
    epochs = 5
    sample_size = 1000  # 设为1000000可快速测试
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    logging.info(f"使用设备: {device}")

    # 数据路径
    train_path = '../kaggle-display-advertising-challenge-dataset/train.txt'

    if not os.path.exists(train_path):
        raise FileNotFoundError(f"数据文件不存在: {train_path}")

    # 加载训练集
    train_dataset = DACDataset(
        file_path=train_path,
        sample_size=sample_size
    )

    # 划分训练集和验证集
    train_indices, val_indices = train_test_split(
        range(len(train_dataset)),
        test_size=0.2,
        random_state=42,
        stratify=train_dataset.data[train_dataset.target_col_index]
    )
    train_subset = Subset(train_dataset, train_indices)
    val_subset = Subset(train_dataset, val_indices)

    # 创建数据加载器
    train_loader = DataLoader(
        train_subset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=4,
        pin_memory=False,
        drop_last=True
    )

    val_loader = DataLoader(
        val_subset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=4,
        pin_memory=False
    )

    # 初始化模型
    model = DeepFM(
        # 离散纬度
        cat_dims=train_dataset.cat_dims,
        # 连续纬度
        num_cols_count=len(train_dataset.num_cols_indices),
        embedding_dim=embedding_dim,
        hidden_dims=hidden_dims
    )

    # 训练配置
    criterion = nn.BCELoss()
    optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=1e-5)

    # 开始训练
    logging.info("开始训练DeepFM模型...")
    trained_model = train_model(
        model, train_loader, val_loader,
        criterion, optimizer, device,
        epochs=epochs
    )

    # 保存最终模型
    torch.save({
        'model_state_dict': trained_model.state_dict(),
        'cat_encoders': train_dataset.cat_encoders,
        'num_scaler': train_dataset.num_scaler
    }, '../res/deepfm_dac_v1_final.pth')
    logging.info("训练完成，模型已保存为 deepfm_dac_v1_final.pth")


if __name__ == "__main__":
    main()
